package za.org.jcicada.katcp;

/**
 * BaseServer.java
 *
 * @author Simon Cross
 *
 * Created on October 13, 2008, 13:45
 */

import java.io.IOException;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.net.InetSocketAddress;
import java.net.ServerSocket;
import java.nio.ByteBuffer;
import java.nio.CharBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.nio.charset.Charset;
import java.nio.charset.CharsetDecoder;
import java.nio.charset.CharsetEncoder;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

/**
 * Represents a KATCP device.
 *
 * The standard request methods are implemented in katcp.Server
 * and devices which intend to implement the standard request methods
 * should inherit from that instead. Only devices which do not intend
 * to implement the stabdard request methods should inherit directly
 * from BaseServer.
 */
public class BaseServer implements Runnable {

    private static Log logger = LogFactory.getLog(BaseServer.class);

    /**
     * Decoder for decoding incoming message strings.
     */
    protected static CharsetDecoder wireDecoder = Charset.forName("ascii").newDecoder();

    /**
     * Encoder for encoding outgoing message strings.
     */
    protected static CharsetEncoder wireEncoder = Charset.forName("ascii").newEncoder();

    /**
     * Map from handler names to request handlers.
     */
    protected final Map<String, IRequestHandler> requestHandlers
        = new HashMap<String, IRequestHandler>();

    /**
     * Map from sensor names to sensors.
     */
    protected final Map<String, ISensor> sensors
        = new HashMap<String, ISensor>();

    /**
     * Whether the server is ready to run (i.e. not in the process of stopping).
     */
    // canRun is kept true except when the server is stopping
    protected boolean canRun = true;

    /**
     * Flag that determines whether the server is currently running.
     */
    protected boolean running = false;

    /**
     * Parser instance for parsing strings into Message objects.
     */
    protected final Parser parser = new Parser();

    /**
     * Host to listen on.
     */
    protected final String host;

    /**
     * Port to listen on.
     */
    protected final int port;

    /**
     * Socket the server is listening on.
     */
    // server socket is an object attribute so subclass can get to it
    protected ServerSocket server;

    /**
     * Socket channel to the current client.
     */
    protected SocketChannel clientChannel;

    /**
     * MessageSendException (so we can distinguish our IOErrors
     * from the errors of others).
     */
    public static class MessageSendException extends IOException {
        private static final long serialVersionUID = 1L;

        /**
         * Constructs a MessageSendException object.
         *
         * @param message  Reason for the exception.
         */
        public MessageSendException (String message) {
            super(message);
        }

        /**
         * Constructs a MessageSendException object.
         *
         * @param message  Reason for the exception.
         * @param cause  The exception that caused the MessageSendException.
         */
        public MessageSendException (String message, Throwable cause) {
            super(message);
            initCause(cause);
        }
    }

    /**
     * Create a BaseServer.
     *
     * @param host  Host to listen on.
     * @param port  Port to listen on.
     */
    public BaseServer(String host, int port) {
        this.host = host;
        this.port = port;
    }

    /**
     * Register a request handler on the server.
     *
     * @param r  Request handler to register.
     */
    public void addRequestHandler(IRequestHandler r) {
        requestHandlers.put(r.getName(), r);
    }

    /**
     * Register a sensor with the server.
     *
     * @param s  Sensor to register.
     */
    public void addSensor(ISensor s) {
        sensors.put(s.getName(), s);
    }

    /**
     * Create a katcp logging inform message.
     *
     * Usually this will be called from inside a DeviceLogger object,
     * but it is also used by the methods in this class when errors
     * need to be reported to the client.
     *
     * @param levelName Name of level to log at.
     * @param logMsg Log message.
     * @param loggerName Name of the logger to log on.
     * @return An inform log Message.
     */
    public Message logMessage(String levelName, String logMsg,
                              String loggerName) {
        return Message.inform("log",
                levelName,
                String.valueOf(new Date().getTime()), // time since epoch in ms
                loggerName,
                logMsg
        );
    }

    /**
     * Send a message to a client channel.
     *
     * @param sc The channel to the client.
     * @param msg The Message to send.
     * @throws MessageSendException If the attempt to send the message fails.
     */
    public void sendMessage(SocketChannel sc, Message msg)
    throws MessageSendException {
        if (sc == null) {
            throw new MessageSendException("Attempt to send message to non-existent client.");
        }
        try {
            ByteBuffer buf = wireEncoder.encode(CharBuffer.wrap(msg.toString() + "\n"));
            sc.write(buf);
        }
        catch (IOException e) {
            throw new MessageSendException("Failed to send message.", e);
        }
    }

    /**
     * Send an inform message to a client channel.
     *
     * @param sc The channel to the client.
     * @param name The name of the inform Message.
     * @param arguments The inform Message arguments.
     * @throws MessageSendException If the attempt to send the message fails.
     */
    public void inform(SocketChannel sc, String name, String... arguments)
    throws MessageSendException {
        sendMessage(sc, new Message(Message.Type.INFORM, name, arguments)); 
    }

    /**
     * Send an inform message to a client channel.
     *
     * @param sc The channel to the client.
     * @param msg The inform Message to send.
     * @throws MessageSendException If the attempt to send the message fails.
     */
    public void inform(SocketChannel sc, Message msg)
    throws MessageSendException {
        sendMessage(sc, msg);
    }

    /**
     * Send an inform message to all clients.
     *
     * @param name The name of the inform Message.
     * @param arguments The inform Message arguments.
     * @throws MessageSendException If the attempt to send the message fails.
     */
    public void informAll(String name, String... arguments)
    throws MessageSendException {
        sendMessage(clientChannel, new Message(Message.Type.INFORM, name, arguments)); 
    }

    /**
     * Send an inform message to all clients.
     *
     * @param msg The inform Message to send.
     * @throws MessageSendException If the attempt to send the message fails.
     */
    public void informAll(Message msg)
    throws MessageSendException {
        sendMessage(clientChannel, msg); 
    }

    /**
     * Handle a message of any type.
     *
     * @param sc The channel to the client.
     * @param msg The Message to process.
     * @throws MessageSendException If the attempt to process the message fails.
     */
    public void handleMessage(SocketChannel sc, Message msg)
    throws MessageSendException {
        Message.Type msgType = msg.getType();
        
        if (msgType == Message.Type.REQUEST)
            handleRequest(sc, msg);
        else if (msgType == Message.Type.INFORM)
            handleInform(sc, msg);
        else if (msgType == Message.Type.REPLY)
            handleReply(sc, msg);
        else {
            // this should never happen (unless someone adds more types)
            String reason = "Server encountered message which wasn't a"
                          + " request, reply or inform messaage.";
            sendMessage(sc, logMessage("error", reason, "root"));
        }
    }

    /**
     * Handle a request message.
     *
     * @param sc The channel to the client.
     * @param msg The request Message to process.
     * @throws MessageSendException If the attempt to process the message fails.
     */
    public void handleRequest(SocketChannel sc, Message msg)
    throws MessageSendException {
        String msgName = msg.getName();
        IRequestHandler h = requestHandlers.get(msgName);
        Message returnMsg;

        if (h == null) {
            returnMsg = Message.reply(msgName, "invalid",
                "Unknown request.");
        }
        else {

            try {
                returnMsg = h.handle(sc, msg);
            }
            catch (Exception e) {
                StringWriter sw = new StringWriter();
                e.printStackTrace(new PrintWriter(sw));
                returnMsg = Message.reply(msgName, "fail", sw.toString());
            }

        }

        sendMessage(sc, returnMsg);
    }

    /**
     * Handle an inform message.
     *
     * This base class simply ignores inform messages.
     *
     * @param sc The channel to the client.
     * @param msg The inform Message to process.
     * @throws MessageSendException If the attempt to process the message fails.
     */
    public void handleInform(SocketChannel sc, Message msg)
    throws MessageSendException { 
    }

    /**
     * Handle a reply message.
     *
     * This base class sends an inform #log error message
     * on receipt of a reply since devices should not receive
     * replies from standard messages.
     *
     * @param sc The channel to the client.
     * @param msg The reply Message to process.
     * @throws MessageSendException If the attempt to process the message fails.
     */
    public void handleReply(SocketChannel sc, Message msg)
    throws MessageSendException {
        String reason = "Unexpected reply message !" + msg.getName()
                      + " received by server.";
        sendMessage(sc, logMessage("error", reason, "root"));
    }

    /**
     * Handle a chunk of data.
     *
     * @param sc The channel to the client.
     * @param chunk A piece of data from the client to process.
     * @return The unprocessed part of the data.
     * @throws MessageSendException If the attempt to process the chunk fails.
     */
    public String handleChunk(SocketChannel sc, String chunk)
    throws MessageSendException {
        chunk = chunk.replace('\r', '\n');
        String [] lines = chunk.split("\\n", -1);
        String fullLine;
        Message msg;

        for (int i=0; i < lines.length-1; i++) {
            fullLine = lines[i];
            if (fullLine.length() == 0)
                continue;
            try {
                msg = parser.parse(fullLine);
                handleMessage(sc, msg);
            }
            catch (Exception e) {
                // usually the parser will throw Parser.KatcpSyntaxExceptions
                // but we really do want to catch other errors too.
                StringWriter sw = new StringWriter();
                e.printStackTrace(new PrintWriter(sw));
                sendMessage(sc, logMessage("error", sw.toString(), "root"));
            }
        }

        return lines[lines.length-1];        
    }

    /**
     * Listen for clients and process their requests.
     */
    public void run() {
        ServerSocketChannel serverChannel;
        SocketChannel oldClientChannel = null;
        SocketChannel tmpClientChannel = null;
        Selector selector;
        ByteBuffer buffer = ByteBuffer.allocate(4096);
        CharBuffer cbuffer = CharBuffer.allocate(4096);
        String chunk = "";
        String reason;
        long timeout = 50; // select timeout in ms
        clientChannel = null;

        try {
            selector = Selector.open();
            serverChannel = ServerSocketChannel.open();
            serverChannel.configureBlocking(false);
            server = serverChannel.socket();
            server.bind(new InetSocketAddress(host, port));
            serverChannel.register(selector, SelectionKey.OP_ACCEPT);
            server.setReuseAddress(true);
        }
        catch (IOException e) {
            throw new RuntimeException("Failed to start device server", e);
        }

        synchronized (this) {
            running = canRun;
        }

        while(running) {
            try {
                selector.select(timeout);
            }
            catch (IOException e) {
                throw new RuntimeException("Server selector threw an IOException", e);
            }

            for(SelectionKey key: selector.selectedKeys()) {
                if ((key.readyOps() & SelectionKey.OP_ACCEPT) != 0) {
                    oldClientChannel = clientChannel;

                    try {
                        clientChannel = serverChannel.accept();
                        if (clientChannel != null) {
                            clientChannel.configureBlocking(false);
                            clientChannel.register(selector, SelectionKey.OP_READ);
                        }
                    }
                    catch (IOException e) {
                        if (clientChannel != null) {
                            try { clientChannel.close(); }
                            catch (IOException closeExp) {}
                        }
                        clientChannel = null;
                    }

                    // bail out if new client setup failed
                    if (clientChannel == null) {
                        clientChannel = oldClientChannel;
                        continue;
                    }
                    
                    // clean up old client if needed
                    if (oldClientChannel != null) {
                        reason = "New client connected from "
                               + clientChannel.socket().getInetAddress().toString();

                        try { onClientDisconnect(oldClientChannel, reason); }
                        catch (MessageSendException e) {}

                        try { oldClientChannel.close(); }
                        catch (IOException e) {}
                    }

                    // setup new client
                    if (clientChannel != null) {
                        try { onClientConnect(clientChannel); }
                        catch (MessageSendException e) {
                            try { clientChannel.close (); }
                            catch (IOException closeExp) {}
                            clientChannel = null;
                        }
                    }

                    chunk = "";
                }
                else if ((key.readyOps() & SelectionKey.OP_READ) != 0) {
                    tmpClientChannel = (SocketChannel) key.channel();

                    buffer.clear();
                    try {
                        tmpClientChannel.read(buffer);
                    }
                    catch (IOException e) {
                        try { tmpClientChannel.close(); }
                        catch (IOException closeExp) {}
                        continue;
                    }

                    buffer.flip();
                    if (buffer.limit() == 0) {
                        // read flag with no data == client closed
                        try { tmpClientChannel.close(); }
                        catch (IOException closeExp) {}
                        continue;
                    }

                    try {
                        cbuffer.clear();
                        wireDecoder.decode(buffer, cbuffer, true);
                        cbuffer.flip();
                        chunk = handleChunk(tmpClientChannel, chunk + cbuffer.toString());
                    }
                    catch (MessageSendException e) {
                        logger.error(String.format("Error handling chunk from client ['%s']", String.valueOf(chunk)));
                        chunk = ""; // discard any partial chunks
                    }
                }
            }
            
            selector.selectedKeys().clear();
        }

        if (clientChannel != null) {
            reason = "Device server shutting down.";
 
            try { onClientDisconnect(clientChannel, reason); }
            catch (MessageSendException e) {}
 
            try { clientChannel.close(); }
            catch (IOException e) {}
        }

        try {
            serverChannel.close();
        } catch (IOException e) {
                StringWriter sw = new StringWriter();
                e.printStackTrace(new PrintWriter(sw));
                logger.error("Device server failed to shutdown with error: " + sw.toString());
        }

        // prepare running for next call to run()
        synchronized (this) {
            canRun = true;
        }
    }


    /**
     * Return port this server is listening on or -1 if
     * it is not running.
     *
     * @return The port the server is listening on or -1 if it is not running.
     */
    public synchronized int getLocalPort() {
        if (running)
            return server.getLocalPort();
        else
            return -1;
    }

    /**
     * Check whether a server is running.
     *
     * @return True if the server is running, false otherwise.
     */
    public synchronized boolean isRunning() {
        return running;
    }

    /**
     * Stop a running server.
     */
    public synchronized void stop() {
        running = false;
        canRun = false;
    }

    /**
     * Called after client connection is established.
     *
     * Subclasses should override if they wish to send clients
     * message or perform house-keeping at this point.
     *
     * @param sc The channel to the client.
     * @throws MessageSendException If an attempt to communicate with the client fails.
     */
    public void onClientConnect(SocketChannel sc)
    throws MessageSendException {
    }

    /**
     * Called before a client connection is closed.
     *
     * Subclasses should override if they wish to send clients
     * message or perform house-keeping at this point. The server
     * cannot guarantee this will be called (for example, the client
     * might drop the connection). The message parameter contains
     * the reason for the disconnection.
     *
     * @param sc The channel to the client.
     * @param reason The reason the client is being disconnected.
     * @throws MessageSendException If an attempt to communicate with the client fails.
     */
    public void onClientDisconnect(SocketChannel sc, String reason)
    throws MessageSendException {
    }

}
